#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/4/4 13:19
# @Author  : Grayson Liu
# @Email   : graysonliu@foxmail.com
# @File    : GeoGraphicProcess.py


from osgeo import gdal
import numpy as np
import cv2
import json
import re
import math
from scipy import ndimage

from CustomSignal import CustomSignal


class GeoGraphicProcess:
    global_signal = CustomSignal()

    def __init__(self):
        """
        初始化，读取WorldView遥感影像数据
        """
        # 加载方法的配置文件method.json
        # 相关信息存储在method.json文件中，是字典
        self.method_data = json.load(open("config/method.json", "r", encoding="utf-8"))

        # 成员变量
        self.x_size: int = 0  # 图像长度
        self.y_size: int = 0  # 图像宽度
        self.raster_count: int = 0  # 获取波段数
        self.projection: str = ""  # 获取投影，返回投影字符串
        self.transform: tuple = ()  # 获取仿射矩阵信息，返回一个元组
        self.im_data: np.ndarray = None  # 所有波段的图像矩阵
        self.current_data: np.ndarray = None  # 当前需要处理的图像矩阵

    def read_file(self, path: str):
        """
        读取WorldView遥感影像数据
        :param path: 遥感影像数据文件绝对路径
        :return:
        """
        # 读取图像的基本信息和地理信息
        dataset = gdal.Open(path)  # 打开遥感影像，返回gdal.Dataset类型的对象
        self.x_size = dataset.RasterXSize  # 图像长度
        self.y_size = dataset.RasterYSize  # 图像宽度
        self.raster_count = dataset.RasterCount  # 获取波段数
        self.projection = dataset.GetProjection()  # 获取投影，返回投影字符串
        self.transform = dataset.GetGeoTransform()  # 获取仿射矩阵信息，返回一个元组

        # 读取图像矩阵，保存在np.ndarray中
        self.im_data = dataset.ReadAsArray(0, 0, self.x_size, self.y_size)

    def get_graphic(self, raster_band: int) -> np.ndarray:
        """
        获取遥感影像某一个波段的图像矩阵
        :param raster_band: 波段号，范围从1到8
        :return: np.ndarray类型的图像矩阵
        """
        return self.im_data[raster_band - 1, :, :]

    def get_geo_info(self) -> dict:
        """
        获取遥感影像基本信息和地理信息，用字典返回
        :return: 包含信息的字典，键值对
        """
        geo_info = {"X size": str(self.x_size),
                    "Y size": str(self.y_size),
                    "Raster Count": str(self.raster_count),
                    "Projection": self.projection,
                    "Transform": str(self.transform)
                    }
        return geo_info

    def save_graphic(self, path, quality, raster_band):
        """
        保存遥感影像数据，保存为jpg格式文件，便于查看
        :param path: 保存文件路径
        :param quality: 品质，范围从0到100
        :param raster_band: 波段号，范围从1到8
        :return: None
        """
        img = self.im_data[raster_band - 1, :, :]
        # 归一化，灰度值范围变为0到255
        cv2.normalize(src=img, dst=img, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)
        cv2.imwrite(path, img, [int(cv2.IMWRITE_JPEG_QUALITY), quality])

    def save_current_graphic(self, path, quality):
        """
        保存当前current_data的处理结果，jpg格式
        :param path: 保存文件路径
        :param quality: 品质
        :return:
        """
        cv2.imwrite(path, self.current_data, [int(cv2.IMWRITE_JPEG_QUALITY), quality])

    def set_input_graphic(self, input_dic: dict):
        """
        根据传入的MethodList中的”input“字典，选择相应波段或指数计算方法，生成后续需要处理的图像矩阵self.current_data
        :param input_dic: MethodList中的”input“字典
        :return:
        """

        def convert(matched):
            """
            转换函数，将识别的 {X} 转换为 [X-1,:,:] 的形式，传递给正则表达式替换re.sub()使用
            :param matched:
            :return:
            """
            index = int(matched.group("value")) - 1
            return "self.im_data[{}, :, :]".format(index)

        # 如果是选择多个频段计算指数
        if input_dic["calculate"] == "multiple":
            # 将类似于“({7}-{5})/({7}+{5})”形式的字符串，把{X}的部分替换为self.im_data[{X-1}, :, :]，X是波段数下标从1开始
            # 这样就可以使用eval()进行表达式求值了
            eval_str = re.sub("{(?P<value>\d+)}", convert, input_dic["raster"])
            self.current_data = eval(eval_str)
        # 如果仅读取单个频段
        else:
            raster_index = int(input_dic["raster"])
            self.current_data = self.im_data[raster_index - 1, :, :]

        # 显示处理信息
        GeoGraphicProcess.global_signal.update_text.emit("max: {}, min: {}, mean: {}".format(
            self.current_data.max(),
            self.current_data.min(),
            self.current_data.mean()
        ))

    def run_preprocess(self, preprocess_dic: dict):
        """
        根据传入的MethodList中的”proprocess“字典，执行缺失值处理和异常值处理的方法，同时将图像数据转为uint8型
        :param preprocess_dic:
        :return:
        """
        # TODO: 这边偷懒，无论选啥都是执行这一种方法，后面要改的话多加几个分支处理即可
        # 简单的异常值筛除，计算均值和标准差，超出均值正负3倍标准差的的样本点认定为异常，用边界值替代
        GeoGraphicProcess.global_signal.update_text.emit("执行缺失值处理方法：{}".format(preprocess_dic["nan_process"]))
        self.current_data[np.isinf(self.current_data)] = 0

        GeoGraphicProcess.global_signal.update_text.emit("执行异常值处理方法：{}".format(preprocess_dic["anomaly_process"]))
        mean = np.nanmean(self.current_data)
        std = np.nanstd(self.current_data)
        upper_limit = mean + 3 * std
        lower_limit = mean - 3 * std
        self.current_data[self.current_data > upper_limit] = upper_limit
        self.current_data[self.current_data < lower_limit] = lower_limit
        self.current_data[np.isnan(self.current_data)] = mean

        # 先进行类型转换，转换为8位int型
        GeoGraphicProcess.global_signal.update_text.emit("归一化处理")
        self.current_data = cv2.normalize(src=self.current_data,
                                          dst=None,
                                          alpha=255,
                                          beta=0,
                                          norm_type=cv2.NORM_MINMAX)
        self.current_data = self.current_data.astype("uint8")

        # 显示处理信息
        GeoGraphicProcess.global_signal.update_text.emit("max: {}, min: {}, mean: {}".format(
            self.current_data.max(),
            self.current_data.min(),
            self.current_data.mean()
        ))

    def run_methods(self, method_num: int, methods: list):
        """
        根据传入的MethodList中的”methods“列表，按照顺序运行每一个方法，方法都是对成员变量self.current_data进行处理
        :param method_num: 方法数量
        :param methods: 方法列表，包含若干字典
        :return:
        """
        if method_num == 0:
            return
        for i in range(method_num):
            method = methods[i]
            GeoGraphicProcess.global_signal.update_text.emit("<b>步骤 {}/{}</b> 正在处理方法：{}".format(i + 1,
                                                                                                method_num,
                                                                                                method["name"]))
            method_dic = self.method_data[method["type_index"]][method["method_index"]]
            function_name = method_dic["function"]
            if method_dic["parameter_number"] == 0:
                parameters_str = ""
            else:
                parameter_list = [str(i["value"]) for i in method["parameters"]]
                parameters_str = ", ".join(parameter_list)
            # 根据函数名和参数列表拼接成语句，使用exec运行相应的方法
            exec("self.{}({})".format(function_name, parameters_str))
            # 显示处理信息
            GeoGraphicProcess.global_signal.update_text.emit("max: {}, min: {}, mean: {:.4f}".format(
                self.current_data.max(),
                self.current_data.min(),
                self.current_data.mean()
            ))

    def run_method_single(self, method: dict, method_num: int, current_index: int):
        """
        执行单步方法，传入一个方法字典，对成员变量self.current_data进行处理
        :param method: 一个方法字典
        :param method_num: 方法总数
        :param current_index: 当前方法在方法列表里面的下标
        :return:
        """
        GeoGraphicProcess.global_signal.update_text.emit("<b>步骤 {}/{}</b> 正在处理方法：{}".format(current_index + 1,
                                                                                            method_num,
                                                                                            method["name"]))
        method_dic = self.method_data[method["type_index"]][method["method_index"]]
        function_name = method_dic["function"]
        if method_dic["parameter_number"] == 0:
            parameters_str = ""
        else:
            parameter_list = [str(i["value"]) for i in method["parameters"]]
            parameters_str = ", ".join(parameter_list)
        # 根据函数名和参数列表拼接成语句，使用exec运行相应的方法
        exec("self.{}({})".format(function_name, parameters_str))
        # 显示处理信息
        GeoGraphicProcess.global_signal.update_text.emit("max: {}, min: {}, mean: {:.4f}".format(
            self.current_data.max(),
            self.current_data.min(),
            self.current_data.mean()
        ))

    # 以下函数命名均不可更改
    def hist_equalization(self):
        """
        直方图均衡化
        :return:
        """
        self.current_data = cv2.equalizeHist(self.current_data)

    def hist_specification(self):
        """
        直方图规定化，使用局部自适应直方图匹配进行
        :return:
        """
        # 局部自适应直方图匹配（用于增强对比度）
        clahe = cv2.createCLAHE(clipLimit=0.01, tileGridSize=(100, 100))
        self.current_data = clahe.apply(self.current_data)

    def average_blur(self, kernel_size: int):
        """
        均值滤波
        :param kernel_size: 模板大小
        :return:
        """
        self.current_data = cv2.blur(self.current_data, (kernel_size, kernel_size))

    def median_blur(self, kernel_size: int):
        """
        中值滤波
        :param kernel_size: 模板大小
        :return:
        """
        self.current_data = cv2.medianBlur(self.current_data, kernel_size)

    def grad(self):
        """
        梯度算子锐化
        :return:
        """
        pass

    def laplacian(self):
        """
        Laplacian算子锐化
        :return:
        """
        self.current_data = cv2.Laplacian(self.current_data, cv2.CV_8U, ksize=3)

    def sobel(self):
        """
        Sobel算子锐化
        :return:
        """
        x_sobel = cv2.Sobel(self.current_data, cv2.CV_16S, 1, 0)
        y_sobel = cv2.Sobel(self.current_data, cv2.CV_16S, 0, 1)

        abs_x = cv2.convertScaleAbs(x_sobel)
        abs_y = cv2.convertScaleAbs(y_sobel)

        result = cv2.addWeighted(abs_x, 0.5, abs_y, 0.5, 0)
        self.current_data = result

    def prewitt(self):
        """
        Prewitt算子锐化
        :return:
        """
        prewitt_1 = np.array([[-1, 0, 1],
                              [-1, 0, 1],
                              [-1, 0, 1]])

        prewitt_2 = np.array([[-1, -1, -1],
                              [0, 0, 0],
                              [1, 1, 1]])
        x_prewitt = cv2.filter2D(self.current_data, -1, prewitt_1)
        y_prewitt = cv2.filter2D(self.current_data, -1, prewitt_2)
        result_data = np.zeros(self.current_data.shape)
        for i in range(self.current_data.shape[0]):
            for j in range(self.current_data.shape[1]):
                result_data = max(x_prewitt[i][j], y_prewitt[i][j])
        self.current_data = result_data

    def kirsch(self):
        """
        Kirsch算子锐化
        :return:
        """
        pass

    def isotropic(self):
        """
        Isotropic算子锐化
        :return:
        """
        factor = math.sqrt(2)
        isotropic_1 = np.array([[-1, 0, 1],
                              [-factor, 0, factor],
                              [-1, 0, 1]])

        isotropic_2 = np.array([[-1, -factor, -1],
                              [0, 0, 0],
                              [1, factor, 1]])
        x_isotropic = cv2.filter2D(self.current_data, -1, isotropic_1)
        y_isotropic = cv2.filter2D(self.current_data, -1, isotropic_2)
        result_data = np.zeros(self.current_data.shape)
        for i in range(self.current_data.shape[0]):
            for j in range(self.current_data.shape[1]):
                result_data = max(x_isotropic[i][j], y_isotropic[i][j])
        self.current_data = result_data

    def butterworth_low(self):
        """
        巴特沃斯低通滤波，用高斯模糊代替
        :return:
        """
        self.current_data = cv2.GaussianBlur(self.current_data, (5, 5), 0)

    def butterworth_high(self):
        """
        巴特沃斯高通滤波
        :return:
        """
        def butterworth_low_pass_kernel(img, cut_off, butterworth_order=1):
            r, c = img[1], img[0]
            u = np.arange(r)
            v = np.arange(c)
            u, v = np.meshgrid(u, v)
            low_pass = np.sqrt((u - r / 2) ** 2 + (v - c / 2) ** 2)
            denom = 1.0 + (low_pass / cut_off) ** (2 * butterworth_order)
            low_pass = 1.0 / denom
            return low_pass

        D0 = 5
        n = 1
        kernel = 1 - butterworth_low_pass_kernel(self.current_data, D0, n)
        gray = np.float64(self.current_data)
        gray_fft = np.fft.fft2(gray)
        gray_fftshift = np.fft.fftshift(gray_fft)
        dst = np.zeros_like(gray_fftshift)
        dst_filtered = kernel * gray_fftshift
        dst_ifftshift = np.fft.ifftshift(dst_filtered)
        dst_ifft = np.fft.ifft2(dst_ifftshift)
        dst = np.abs(np.real(dst_ifft))
        dst = np.clip(dst, 0, 255)
        self.current_data = np.uint8(dst)

    def DESTIN_edge_strength(self):
        """
        (DESTIN方法)8-邻域边缘强度计算
        :return:
        """
        x_size = self.x_size
        y_size = self.y_size
        data: np.ndarray = self.current_data.astype("float64")

        strength = np.zeros((x_size, y_size))  # 用于存放边缘强度计算值
        factor = math.sqrt(2) / 2  # 斜方向计算强度时需要的系数

        # 4.1. 四个顶角边缘强度值计算
        strength[0, 0] = abs(data[0, 0] - data[1, 0]) \
                         + abs(factor * (data[0, 0] - data[1, 1])) \
                         + abs(data[0, 0] - data[0, 1])
        strength[0, y_size - 1] = abs(data[0, y_size - 1] - data[1, y_size - 1]) \
                                  + abs(factor * (data[0, y_size - 1] - data[1, y_size - 2])) \
                                  + abs(data[0, y_size - 1] - data[0, y_size - 2])
        strength[x_size - 1, 0] = abs(data[x_size - 1, 0] - data[x_size - 2, 0]) \
                                  + abs(factor * (data[x_size - 1, 0] - data[x_size - 2, 1])) \
                                  + abs(data[x_size - 1, 0] - data[x_size - 1, 1])
        strength[x_size - 1, y_size - 1] = abs(data[x_size - 1, y_size - 1] - data[x_size - 2, y_size - 1]) \
                                           + abs(
            factor * (data[x_size - 1, y_size - 1] - data[x_size - 2, y_size - 2])) \
                                           + abs(
            data[x_size - 1, y_size - 1] - data[x_size - 1, y_size - 2])

        # 4.2. 四条边缘（除四个顶点）边缘强度值计算
        for j in range(1, y_size - 1):
            strength[0, j] = abs(data[0, j] - data[0, j - 1]) \
                             + abs(factor * (data[0, j] - data[1, j - 1])) \
                             + abs(data[0, j] - data[1, j]) \
                             + abs(factor * (data[0, j] - data[1, j + 1])) \
                             + abs(data[0, j] - data[0, j + 1])
            strength[x_size - 1, j] = abs(data[x_size - 1, j] - data[x_size - 1, j - 1]) \
                                      + abs(factor * (data[x_size - 1, j] - data[x_size - 2, j - 1])) \
                                      + abs(data[x_size - 1, j] - data[x_size - 2, j]) \
                                      + abs(factor * (data[x_size - 1, j] - data[x_size - 2, j + 1])) \
                                      + abs(data[x_size - 1, j] - data[x_size - 1, j + 1])
        for i in range(1, x_size - 1):
            strength[i, 0] = abs(data[i, 0] - data[i, 1]) \
                             + abs(factor * (data[i, 0] - data[i - 1, 1])) \
                             + abs(data[i, 0] - data[i - 1, 1]) \
                             + abs(factor * (data[i, 0] - data[i + 1, 1])) \
                             + abs(data[i, 0] - data[i + 1, 0])
            strength[i, y_size - 1] = abs(data[i, y_size - 1] - data[i, y_size - 2]) \
                                      + abs(factor * (data[i, y_size - 1] - data[i - 1, y_size - 2])) \
                                      + abs(data[i, y_size - 1] - data[i - 1, y_size - 1]) \
                                      + abs(factor * (data[i, y_size - 1] - data[i + 1, y_size - 2])) \
                                      + abs(data[i, y_size - 1] - data[i + 1, y_size - 1])

        # 4.3 主体部分边缘强度计算
        for i in range(1, x_size - 1):
            for j in range(1, y_size - 1):
                strength[i, j] = abs(data[i, j] - data[i + 1, j]) \
                                 + abs(data[i, j] - data[i - 1, j]) \
                                 + abs(data[i, j] - data[i, j + 1]) \
                                 + abs(data[i, j] - data[i, j - 1]) \
                                 + abs(factor * (data[i, j] - data[i + 1, j + 1])) \
                                 + abs(factor * (data[i, j] - data[i - 1, j + 1])) \
                                 + abs(factor * (data[i, j] - data[i + 1, j - 1])) \
                                 + abs(factor * (data[i, j] - data[i - 1, j - 1]))

        self.current_data = strength.astype("uint8")  # 数据类型转换，因为中值滤波需要32位浮点数

    def simple_threshold(self, threshold: int):
        """
        简单直方图阈值分割
        :param threshold: 阈值
        :return:
        """
        self.current_data = cv2.threshold(src=self.current_data,
                                          thresh=threshold,
                                          maxval=255,
                                          type=cv2.THRESH_BINARY)

    def adaptive_threshold(self):
        """
        最佳阈值分割，使用局部自适应阈值分割
        :return:
        """
        # 局部自适应阈值分割
        self.current_data = cv2.adaptiveThreshold(self.current_data,
                                                  maxValue=255,
                                                  adaptiveMethod=cv2.ADAPTIVE_THRESH_MEAN_C,
                                                  thresholdType=cv2.THRESH_BINARY, blockSize=5, C=0.0)

    def otsu_threshold(self):
        """
        类间方差阈值分割(Otsu)
        :return:
        """
        self.current_data = cv2.GaussianBlur(src=self.current_data,
                                             ksize=(5, 5),
                                             sigmaX=0)
        result, self.current_data = cv2.threshold(src=self.current_data,
                                                  thresh=0,
                                                  maxval=255,
                                                  type=(cv2.THRESH_BINARY + cv2.THRESH_OTSU))

    def entropy_threshold(self):
        """
        最大熵阈值分割
        :return:
        """
        rows, cols = self.current_data.shape
        gray_hist, bins = np.histogram(self.current_data, bins=256, range=[0, 256])
        norm_gray_hist = gray_hist / float(rows * cols)
        zero_cumu_moment = np.zeros([256], dtype="float32")
        for k in range(256):
            if k == 0:
                zero_cumu_moment[k] = norm_gray_hist[k]
            else:
                zero_cumu_moment[k] = zero_cumu_moment[k - 1] + norm_gray_hist[k]
        entropy = np.zeros([256], dtype="float32")
        for k in range(256):
            if k == 0:
                if norm_gray_hist[k] == 0:
                    entropy[k] = 0
                else:
                    entropy[k] = -norm_gray_hist[k] * np.log10(norm_gray_hist[k])
            else:
                if norm_gray_hist[k] == 0:
                    entropy[k] = entropy[k - 1]
                else:
                    entropy[k] = entropy[k - 1] - norm_gray_hist[k] * np.log10(norm_gray_hist[k])
        ft = np.zeros([256], dtype="float32")
        ft1, ft2 = 0., 0.
        total_entropy = entropy[255]
        for k in range(255):
            max_fornt = np.max(norm_gray_hist[:k + 1])
            max_back = np.max(norm_gray_hist[k + 1: 256])
            if max_fornt == 0 or zero_cumu_moment[k] == 0 or max_fornt == 1 or zero_cumu_moment[k] == 1 or total_entropy == 0:
                ft1 = 0
            else:
                ft1 = entropy[k] / total_entropy * (np.log10(zero_cumu_moment[k]) / np.log10(max_fornt))
            if max_back == 0 or 1 - zero_cumu_moment[k] == 0 or max_back == 1 or 1 - zero_cumu_moment[k] == 1:
                ft2 = 0
            else:
                if total_entropy == 0:
                    ft2 = (np.log10(1 - zero_cumu_moment[k]) / np.log10(max_back))
                else:
                    ft2 = (1 - entropy[k] / total_entropy) * (np.log10(1 - zero_cumu_moment[k]) / np.log10(max_back))
            ft[k] = ft1 + ft2
        thresloc = np.where(ft == np.max(ft))
        thresh = thresloc[0][0]
        self.current_data[self.current_data > thresh] = 255
        self.current_data[self.current_data <= thresh] = 0

    def marr(self, sigma: int, kernel_size: int):
        """
        Marr边缘检测算子
        :param sigma:
        :param kernel_size:
        :return:
        """
        self.current_data = cv2.GaussianBlur(self.current_data, (kernel_size, kernel_size), 0)
        self.current_data = cv2.Laplacian(self.current_data, cv2.CV_8U, ksize=kernel_size)
        # TODO: 应当删除sigma参数

    def canny(self, low_threshold: int, ratio: float, kernel_size: int):
        """
        Canny边缘检测算子
        :param low_threshold:
        :param ratio:
        :param kernel_size:
        :return:
        """
        self.current_data = cv2.GaussianBlur(self.current_data, (kernel_size, kernel_size), 0)
        self.current_data = cv2.Canny(self.current_data, low_threshold, low_threshold * ratio)

    def dilation(self, kernel_size: int):
        """
        形态学膨胀
        :param kernel_size:
        :return:
        """
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
        self.current_data = cv2.dilate(self.current_data, kernel, iterations=1)

    def erosion(self, kernel_size: int):
        """
        形态学腐蚀
        :param kernel_size:
        :return:
        """
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
        self.current_data = cv2.erode(self.current_data, kernel, iterations=1)

    def imfill(self):
        """
        IMFILL（补洞）
        :return:
        """
        self.current_data[np.where(self.current_data == 255)] = 1
        # TODO: structure参数应当让用户选择
        self.current_data = ndimage.binary_fill_holes(self.current_data, structure=np.ones((40, 40))).astype("uint8")
        self.current_data[np.where(self.current_data == 1)] = 255
